package io.sunny.platform.backend.utils;

import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;

/**
 * @Author SUNNY
 * @Date 2025/10/23 11:46
 * @PackageName: io.sunny.platform.backend.utils
 * @Description: 轻量级异步任务协调工具(核心设计：多次提交，一次等待)
 * @Version 1.0
 */
public class LatchUtils {

    /**
     * 使用 ThreadLocal 存储当前线程提交的任务列表，保证线程隔离
     */
    private static final ThreadLocal<List<TaskInfo>> TASK_POOL = ThreadLocal.withInitial(LinkedList::new);

    /**
     * 提交异步任务
     * @param executor 指定线程池执行任务
     * @param runnable 任务逻辑
     */
    public static void submitTask(Executor executor, Runnable runnable) {
        TASK_POOL.get().add(new TaskInfo(executor, runnable));
    }

    /**
     * 获取并清空任务列表
     */
    private static List<TaskInfo> popTask() {
        List<TaskInfo> taskInfos = TASK_POOL.get();
        TASK_POOL.remove();
        return taskInfos;
    }

    /**
     * 触发所有任务执行，并同步等待完成
     * @param timeout 最大等待时间
     * @param timeUnit 时间单位
     * @return 是否在超时前全部完成
     */
    public static boolean waitFor(long timeout, TimeUnit timeUnit) {
        List<TaskInfo> taskInfos = popTask();
        if (taskInfos.isEmpty()) return true;

        CountDownLatch latch = new CountDownLatch(taskInfos.size());
        for (TaskInfo taskInfo : taskInfos) {
            taskInfo.executor.execute(() -> {
                try {
                    taskInfo.runnable.run();
                } finally {
                    latch.countDown();
                }
            });
        }

        try {
            return latch.await(timeout, timeUnit);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            return false;
        }
    }

    /**
     * 内部任务封装类
     */
    private static final class TaskInfo {
        private final Executor executor;
        private final Runnable runnable;

        TaskInfo(Executor executor, Runnable runnable) {
            this.executor = executor;
            this.runnable = runnable;
        }
    }
}
